// Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved

#pragma once

#include <algorithm>
#include <cstdlib>
#include <iostream>
#include <sstream>
#include <string>

namespace kpex {
namespace internal {
inline void check_fail_print(std::stringstream &stream)
{}

template <typename Arg, typename... Rest>
inline void check_fail_print(std::stringstream &stream, Arg &&arg, Rest &&...rest)
{
    stream << std::forward<Arg>(arg);
    check_fail_print(stream, rest...);
}

template <typename... Args>
inline void check_fail(std::string func, std::string file, int line, Args &&...args)
{
    std::stringstream stream;
    stream << "KPEX_CHECK fail in " << func << " at " << file << ":" << line << ", ";
    check_fail_print(stream, std::forward<Args>(args)...);
    stream << "\n";
    std::cerr << stream.str();
    abort();
}
}   //namespace internal
}   //namespace kpex

#define KPEX_CHECK(condition, ...)                                                  \
    do {                                                                           \
        if (__builtin_expect(!(condition), 0)) {                                   \
            kpex::internal::check_fail(__func__, __FILE__, __LINE__, __VA_ARGS__); \
        }                                                                          \
    } while (0)

#define KPEX_CHECK_TENSOR_SHAPE(tensor, ...)                                                                    \
    KPEX_CHECK((tensor).sizes() == c10::IntArrayRef({__VA_ARGS__}), "invalid tensor shape: ", (tensor).sizes(), \ 
        ", expect: ", c10::IntArrayRef({__VA_ARGS__}))

#define KPEX_CHECK_TENSORWRAPPER_SHAPE(tensor, ...)                                 \
    KPEX_CHECK(c10::IntArrayRef((tensor).sizes) == c10::IntArrayRef({__VA_ARGS__}), \
        "invalid tensor wrapper shape: ", c10::IntArrayRef((tensor).sizes),        \
        ", expect: ", c10::IntArrayRef({__VA_ARGS__}))